import sys
from PySide2.QtWidgets import QApplication, QFileDialog, QMainWindow
from PySide2.QtGui import QImage, QPixmap
from PySide2.QtCore import *
from Label_Enhancer_Init import *
from Label_Enhancer_Class import *
from Label_Enhancer_UI import Ui_Label_Enhancer

class MainWindow(QMainWindow):
        
    def __init__(self):
        super().__init__()
        self.ui = Ui_Label_Enhancer()
        self.ui.setupUi(self)

        Global_Signals.info_signal.connect(self.ui_print)
        Global_Signals.key_signal.connect(self.cut_pic)
        sys.stdout = self
        sys.stderr = self
        self.paths = load_path()
        for key, value in self.paths.items():
            if value is None:
                self.paths[key] = QFileDialog.getExistingDirectory(self.ui, 'Select %s path.'%key)
        update_path(self.paths)
        self.update_pics()

    def update_pics(self, pic_original=None):
        result = makedata(self.paths['read'], pic_original)
        if not result is None:
            [pic, self.pic_analysis, self.pic_all, self.pic_r, self.pic_g, self.pic_b, self.pic_mask, self.pic_mask_reverse, self.analysis_data, pic_list] = result
            if pic_original is None:
                self.pic_original = pic
                self.pic_list = pic_list
                self.pic_id = 0
                self.update_pic(self.pic_original, self.ui.pic_original)
            self.update_pic(self.pic_analysis, self.ui.pic_analysis)
            self.update_pic(self.pic_all, self.ui.pic_all)
            self.update_pic(self.pic_r, self.ui.pic_r)
            self.update_pic(self.pic_g, self.ui.pic_g)
            self.update_pic(self.pic_b, self.ui.pic_b)
            self.update_pic(self.pic_mask, self.ui.pic_mask)
            self.update_pic(self.pic_mask_reverse, self.ui.pic_mask_reverse)

    def update_pic(self, pic, label_name):
        height, width, bytesPerComponent = pic.shape
        label_name.ratio = min(label_name.height()/height, label_name.width()/width)
        height = int(label_name.ratio * height)
        width = int(label_name.ratio * width)
        pic = cv2.resize(pic, (width, height))
        bytesPerLine = 3 * width
        cv2.cvtColor(pic, cv2.COLOR_BGR2RGB, pic)
        QImg = QImage(pic.data, width, height, bytesPerLine, QImage.Format_RGB888)
        pixmap = QPixmap.fromImage(QImg)
        label_name.setPixmap(pixmap)

    def write(self, signal):
        Global_Signals.info_signal.emit(signal)

    @QtCore.Slot(str)
    def ui_print(self, text):
        if text != '\n':
            self.ui.info.appendPlainText(text)

    def cut_pic(self, command):
        def change_pic():
            self.pic_original = cv2.imread(self.pic_list[self.pic_id])
            self.update_pic(self.pic_original, self.ui.pic_original)
            self.update_pics(self.pic_original)

        if command == 'enter' and self.ui.pic_original.printable == True:
            [x0, y0, x1, y1] = [self.ui.pic_original.x0, self.ui.pic_original.y0, self.ui.pic_original.x1, self.ui.pic_original.y1]
            ratio = self.ui.pic_original.ratio
            x_min = int(min(x0, x1) / ratio)
            x_max = int(max(x0, x1) / ratio)
            y_min = int(min(y0, y1) / ratio)
            y_max = int(max(y0, y1) / ratio)
            self.pic_cut = self.pic_original[y_min:y_max, x_min:x_max]
            self.update_pics(self.pic_cut)
            self.ui.pic_original.printable == False
        if command == 'a':
            if self.pic_id == 0:
                self.pic_id = len(self.pic_list) - 1
            else:
                self.pic_id -= 1
            change_pic()
        if command == 'd':
            if self.pic_id == len(self.pic_list) - 1:
                self.pic_id = 0
            else:
                self.pic_id += 1
            change_pic()

if __name__ == '__main__':
    QtCore.QCoreApplication.setAttribute(QtCore.Qt.AA_ShareOpenGLContexts)
    app = QApplication([])
    UI = MainWindow()
    UI.show()
    app.exec_()
    